// Copyright 2022 jeff.li. and/or its affiliates.

#pragma once

#include <memory>
#include <string>

#include <tbir/pipeline/userdata_cache.h>
#include <tbir/runtime/container.h>

namespace tbir::runtime {

    class UserDataScopedCache {
    public:
        explicit UserDataScopedCache(String scope) {
            iso_name_ = std::move(scope);
            global_cache_ = UserDataCache::instance();
            global_cache_->AddRef(iso_name_);
        }

        virtual ~UserDataScopedCache() {
            global_cache_->DeRef(iso_name_);
        }

        void Set(string_view cls, string_view name, UserDataRef ud) {
            global_cache_->Set(iso_name_, cls, name, std::move(ud));
        }

        UserDataRef Get(string_view cls, string_view name) {
            return global_cache_->Get(iso_name_, cls, name);
        }

        std::vector<UserDataRef> GetAll() {
            return global_cache_->GetAll(iso_name_);
        }

        void Remove(string_view cls, string_view name) {
            global_cache_->Remove(iso_name_, cls, name);
        }

    private:
        UserDataCache *global_cache_;
        String iso_name_;

        friend class TXSession;
    };

}  // namespace tbir::runtime
